package com.smart.thrift.discoveryandbalance.discovery.zookeeper;

import cn.hutool.json.JSONUtil;
import com.google.common.collect.Lists;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryNTimes;

import java.util.List;
import java.util.Random;

public class ThriftZookeeperServerUtils {

    public static final String HOST_NAME_PATTERN = "(\\d+\\.){3}\\d+";

    public static ThriftZookeeperServerNode getServiceNode(String connectString, String zNodePath) throws Exception {
        CuratorFramework zookeeperClient = null;
        try {
            zookeeperClient = CuratorFrameworkFactory.newClient(connectString, new RetryNTimes(3, 1000));
            zookeeperClient.start();
            List<String> nodes = zookeeperClient.getChildren().watched().forPath(zNodePath);
            int nodeIndex = (new Random()).nextInt(nodes.size());
            String nodeData = new String(zookeeperClient.getData().forPath(zNodePath + "/" + nodes.get(nodeIndex)));
            ThriftZookeeperServerNode nodeInfo = JSONUtil.toBean(nodeData, ThriftZookeeperServerNode.class);
            if (null != nodeInfo && nodeInfo.getHost().matches(HOST_NAME_PATTERN) && nodeInfo.getPort() > 0) {
                return nodeInfo;
            }
            return null;
        } finally {
            if (zookeeperClient != null) {
                zookeeperClient.close();
            }
        }
    }

    public static List<ThriftZookeeperServerNode> getAllServiceNodes(String connectString, String zNodePath) throws Exception {
        CuratorFramework zookeeperClient = null;
        try {
            zookeeperClient = CuratorFrameworkFactory.newClient(connectString, new RetryNTimes(3, 1000));
            zookeeperClient.start();
            List<String> nodes = zookeeperClient.getChildren().watched().forPath(zNodePath);
            List<ThriftZookeeperServerNode> res = Lists.newArrayListWithCapacity(nodes.size());
            for (String node : nodes) {
                String nodeData = new String(zookeeperClient.getData().forPath(zNodePath + "/" + node));
                ThriftZookeeperServerNode nodeInfo = JSONUtil.toBean(nodeData, ThriftZookeeperServerNode.class);
                if (null != nodeInfo && nodeInfo.getHost().matches(HOST_NAME_PATTERN) && nodeInfo.getPort() > 0) {
                    res.add(nodeInfo);
                }
            }
            return res;
        } finally {
            if (zookeeperClient != null) {
                zookeeperClient.close();
            }
        }
    }
}
